import os
import argparse
import pickle
import torch
import numpy as np


def _to_numpy(x):
    if isinstance(x, torch.Tensor):
        return x.detach().cpu().numpy()
    return np.asarray(x)


def _to_seq_feat(x):
    x = _to_numpy(x)
    if x.ndim == 3:
        # (N, C, T) -> (N, T, C)
        return np.transpose(x, (0, 2, 1))
    if x.ndim == 4:
        # (N, S, C, B) -> (N, S, C*B)
        n, s, c, b = x.shape
        return x.reshape(n, s, c * b)
    raise ValueError(f"unsupported samples ndim: {x.shape}")


def _load_split(pt_path):
    obj = torch.load(pt_path, map_location="cpu")
    samples = obj.get("samples")
    labels = obj.get("labels")
    if samples is None or labels is None:
        raise RuntimeError(f"missing keys in {pt_path}")
    X = _to_seq_feat(samples)
    y = _to_numpy(labels).reshape(-1)
    y = y.astype(np.int64)
    y = y.reshape(-1, 1, 1)
    return X, y


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--src_dir", required=True, help="folder containing train.pt/val.pt/test.pt")
    p.add_argument("--dst_root", default="data", help="CLMER data root (will write data/<dataset>/<dataset>.pkl)")
    p.add_argument("--dataset", default="deap9class0", help="dataset name used by CLMER (matches --dataset in fmain.py)")
    p.add_argument("--v1_dim", type=int, default=1)
    p.add_argument("--v2_dim", type=int, default=1)
    args = p.parse_args()

    train_pt = os.path.join(args.src_dir, "train.pt")
    val_pt = os.path.join(args.src_dir, "val.pt")
    test_pt = os.path.join(args.src_dir, "test.pt")

    X_tr, y_tr = _load_split(train_pt)
    X_va, y_va = _load_split(val_pt)
    X_te, y_te = _load_split(test_pt)

    n_tr, L, Fp = X_tr.shape
    n_va = X_va.shape[0]
    n_te = X_te.shape[0]

    V1_tr = np.zeros((n_tr, L, args.v1_dim), dtype=np.float32)
    V1_va = np.zeros((n_va, L, args.v1_dim), dtype=np.float32)
    V1_te = np.zeros((n_te, L, args.v1_dim), dtype=np.float32)

    V2_tr = np.zeros((n_tr, L, args.v2_dim), dtype=np.float32)
    V2_va = np.zeros((n_va, L, args.v2_dim), dtype=np.float32)
    V2_te = np.zeros((n_te, L, args.v2_dim), dtype=np.float32)

    data = {
        "train": {"vision1": V1_tr, "vision2": V2_tr, "physio": X_tr, "labels": y_tr},
        "valid": {"vision1": V1_va, "vision2": V2_va, "physio": X_va, "labels": y_va},
        "test":  {"vision1": V1_te, "vision2": V2_te, "physio": X_te, "labels": y_te},
    }

    dst_dir = os.path.join(args.dst_root, args.dataset)
    os.makedirs(dst_dir, exist_ok=True)
    dst_pkl = os.path.join(dst_dir, f"{args.dataset}.pkl")
    with open(dst_pkl, "wb") as f:
        pickle.dump(data, f)

    print(dst_pkl)


if __name__ == "__main__":
    main()
